package org.ws.task.zookeeper;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.lang.StringUtils;
import org.apache.curator.ensemble.fixed.FixedEnsembleProvider;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.ACLProvider;
import org.apache.curator.framework.imps.GzipCompressionProvider;
import org.apache.curator.retry.BoundedExponentialBackoffRetry;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.data.ACL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.event.support.BootstrapEventHolder;
import org.ws.task.event.support.startup.BootstrapEvent;
import org.ws.task.event.support.startup.BootstrapEventType;
import org.ws.task.listener.RegisterListener;
import org.ws.task.listener.zk.NodeListener;
import org.ws.task.model.NodeMetadata;
import org.ws.task.model.NodeStatus;
import org.ws.task.zookeeper.listener.NodeEventListener;
import org.ws.task.zookeeper.listener.NodeListenerManager;
import org.ws.task.zookeeper.node.LeaderNode;
import org.ws.task.zookeeper.node.NodeDataCodec;
import org.ws.task.zookeeper.node.WorkerNode;

import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.UUID;

public class DefaultZkNodeManager implements ZkNodeManager, RegisterListener<NodeListener> {

    private static Logger log = LoggerFactory.getLogger(DefaultZkNodeManager.class.getName());

    private final NodeListenerManager nodeListenerManager;
    private final CuratorFramework curatorClient;
    private final ZkConfig zkConfig;
    private final String id;
    private final LeaderNode leaderNode;
    private final WorkerNode workerNode;

    public DefaultZkNodeManager(ZkConfig zkConfig) {
        this.id = UUID.randomUUID().toString();
        this.zkConfig = zkConfig;
        this.curatorClient = markCurator();

        this.nodeListenerManager = new NodeListenerManager();

        this.leaderNode = new LeaderNode(this.id,this.zkConfig.getNamespace(),this.curatorClient,this.nodeListenerManager);
        this.workerNode = new WorkerNode(this.id,this.zkConfig.getNamespace(),this.curatorClient,this.nodeListenerManager);
        // 发送事件
        this.nodeListenerManager.registerListener(new NodeEventListener());
    }

    @Override
    public void registerListener(NodeListener listener) {
        this.nodeListenerManager.registerListener(listener);
    }

    @Override
    public String id() {
        return this.id;
    }

    @Override
    public NodeMetadata getLeader() {
        return this.leaderNode.getLeader();
    }

    @Override
    public boolean isLeader() {
        return this.leaderNode.isLeader();
    }

    @Override
    public void updateLeader(int port, NodeStatus status) {
        this.leaderNode.updateLeader(port,status);
    }

    @Override
    public void updateWorker(int port, NodeStatus status) {
        this.workerNode.updateWorker(port,status);
    }

    @Override
    public void start() {
        this.curatorClient.start();
        this.workerNode.start();
        this.leaderNode.start();
        // sync all worker nodes
        this.getWorkers().stream().forEach(node ->
            BootstrapEventHolder.getInstance().publishEvent(new BootstrapEvent(node, BootstrapEventType.UPDATE_WORKER))
        );
        log.info("Zk Node Manager start.");
    }

    @Override
    public void shutdown() {
        this.workerNode.shutdown();
        this.leaderNode.shutdown();
        this.curatorClient.close();
    }

    @Override
    public Collection<NodeMetadata> getWorkers() {
        return this.workerNode.getWorkers();
    }

    @Override
    public String getNamespace() {
        return this.zkConfig.getNamespace();
    }

    @Override
    public NodeMetadata getLeaderDataNode() throws Exception {
        String path = NodePaths.getLeaderDataPath(getNamespace());
        byte[] bytes = this.curatorClient.getData().forPath(path);
        return NodeDataCodec.decode(bytes, NodeMetadata.class);
    }

    @Override
    public Collection<Map<String,Object>> getLeaderNodes() throws Exception {
        String leaderPath = NodePaths.getLeaderPath(getNamespace());
        List<String> paths = this.curatorClient.getChildren().forPath(leaderPath);
        List<Map<String,Object>> list = Lists.newArrayList();
        for (String path : paths) {
            Map<String,Object> map = Maps.newHashMap();
            byte[] bytes = this.curatorClient.getData().forPath(leaderPath+"/"+path);
            map.put(path,new String(bytes,StandardCharsets.UTF_8));
            list.add(map);
        }
        return list;
    }

    @Override
    public Collection<NodeMetadata> getWorkerNodes() {
        return this.workerNode.getWorkers();
    }

    private CuratorFramework markCurator(){
        final CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
        // auth
        if(StringUtils.isNotBlank(this.zkConfig.getUser())
                && StringUtils.isNotBlank(this.zkConfig.getPwd())){
            byte [] auth = String.format("%s:%s", this.zkConfig.getUser(), this.zkConfig.getPwd()).getBytes(StandardCharsets.UTF_8);
            builder.authorization(this.zkConfig.getAuthScheme(),auth);
        }
        builder.ensembleProvider(new FixedEnsembleProvider(this.zkConfig.getAddress()))
                .sessionTimeoutMs(this.zkConfig.getSessionTimeout())
                .connectionTimeoutMs(this.zkConfig.getConnectionTimeout())
                .retryPolicy(new BoundedExponentialBackoffRetry(1_000,60_000,1000));
        // compression
        if(this.zkConfig.isEnableCompression()){
            builder.compressionProvider(new GzipCompressionProvider());
        }
        // acl
        if(this.zkConfig.isEnableAcl()){
            builder.aclProvider(new ACLProvider() {
                @Override
                public List<ACL> getDefaultAcl() {
                    return ZooDefs.Ids.CREATOR_ALL_ACL;
                }

                @Override
                public List<ACL> getAclForPath(String s) {
                    return ZooDefs.Ids.CREATOR_ALL_ACL;
                }
            });
        }
        return builder.build();
    }
}
